import json
from fastapi import APIRouter, Query
from fastapi.responses import JSONResponse
from typing import Optional, List

from langchain_community.llms import Ollama
from langchain_community.document_loaders import (
    TextLoader,
    CSVLoader,
    DirectoryLoader,
    UnstructuredHTMLLoader,
    JSONLoader,
    PyPDFLoader
)

router = APIRouter()

# Ollama LLM 实例
ollama = Ollama(base_url='http://172.16.21.38:11436', model="qwen3:0.6b")

@router.get("/test")
async def test():
    return JSONResponse(content={"message": "Hello, World!"})

@router.get("/load_demo")
async def load_demo(file_type: str = Query(..., description="文件类型: markdown/csv/html/json/pdf/dir")):
    """
    演示不同类型文件的 Loader 机制
    """
    docs = []
    if file_type == "markdown":
        loader = TextLoader("H:/DWORPLACETEST/PythonWebTemplate/AIAgentBase/app/testapi/D4RAG/examples/example.md", encoding="utf-8")
        docs = loader.load()

    elif file_type == "csv":
        loader = CSVLoader(file_path="H:/DWORPLACETEST/PythonWebTemplate/AIAgentBase/app/testapi/D4RAG/examples/loader.csv")
        docs = loader.load()

    elif file_type == "dir":
        loader = DirectoryLoader("H:/DWORPLACETEST/PythonWebTemplate/AIAgentBase/app/testapi/D4RAG/examples/", glob="**/*.txt")
        docs = loader.load()

    elif file_type == "html":
        loader = UnstructuredHTMLLoader("H:/DWORPLACETEST/PythonWebTemplate/AIAgentBase/app/testapi/D4RAG/examples/loader.html")
        docs = loader.load()

    elif file_type == "json":
        loader = JSONLoader(
            file_path="examples/example.json",
            jq_schema=".items[]",   # jq语法，提取字段
            text_content=False
        )
        docs = loader.load()

    elif file_type == "pdf":
        loader = PyPDFLoader("H:/DWORPLACETEST/PythonWebTemplate/AIAgentBase/app/testapi/D4RAG/examples/loader.pdf")
        docs = loader.load()

    else:
        return JSONResponse(content={"error": "不支持的文件类型"})

    # 只返回前 200 个字符作为展示
    content_preview = docs[0].page_content[:200] if docs else "空文档"

    return JSONResponse(content={
        "file_type": file_type,
        "docs_count": len(docs),
        "preview": content_preview
    })
